import os
import pickle
import glob

import cv2
import numpy as np

LABELS = [
    "airplane", "automobile", "bird", "cat", "deer",
    "dog", "frog", "horse", "ship", "truck",
]

TRAIN_DATA_PATH = "./datasets/TRAIN"
TEST_DATA_PATH = "./datasets/TEST"


def unpickle(file):
    # import pickle
    with open(file, 'rb') as fo:
        d = pickle.load(fo, encoding='bytes')
    return d


def sava_img(root, im_label, im_filename, im_data):
    """
    保存文件到指定目录
    :param root: 根目录， train / test
    :param im_label: 图片标签,如：airplane
    :param im_filename: 图片文件名
    :param im_data: 数据 32*32*3
    :return:
    """
    _path = os.path.join(root, im_label)
    if not os.path.exists(_path):
        os.makedirs(_path, exist_ok=True)

    cv2.imwrite(os.path.join(_path, im_filename), im_data)


def read_train():
    train_list = glob.glob("./datasets/data_batch_*")
    for file in train_list:
        # 每个文件中存了10000张图片
        item = unpickle(file)
        # print(item.keys()) # dict_keys([b'batch_label', b'labels', b'data', b'filenames'])
        for im_idx, im_data in enumerate(item[b'data']):
            # im_data是以向量存储的，需要reshape。cifar10的数据是通道优先的，所以需要reshape为3*32*32的形式，然后再转化为32*32*3的形式
            # print(im_idx)
            # print(im_data)
            im_label = item[b"labels"][im_idx]
            im_filename = item[b"filenames"][im_idx]
            # print(im_label, im_filename, im_data)

            im_label_name = LABELS[im_label]
            im_filename = im_filename.decode()
            im_data = np.array(im_data).reshape((3, 32, 32)).transpose(1, 2, 0)

            # print(im_label_name, im_filename, im_data)
            # # cv2.imshow("im_data", im_data)
            # cv2.imshow("im_data", cv2.resize(im_data, (200, 200)))
            # cv2.waitKey(0)
            if (im_idx + 1) % 1000 == 0:
                print(f"正在写入{file}数据集中，第{im_idx + 1}张图片：{im_filename}")

            sava_img(TRAIN_DATA_PATH, im_label_name, im_filename, im_data)


def read_test():
    test_list = glob.glob("./datasets/test_batch*")
    for file in test_list:
        # 每个文件中存了10000张图片
        item = unpickle(file)
        # print(item.keys()) # dict_keys([b'batch_label', b'labels', b'data', b'filenames'])
        for im_idx, im_data in enumerate(item[b'data']):
            # im_data是以向量存储的，需要reshape。cifar10的数据是通道优先的，所以需要reshape为3*32*32的形式，然后再转化为32*32*3的形式
            # print(im_idx)
            # print(im_data)
            im_label = item[b"labels"][im_idx]
            im_filename = item[b"filenames"][im_idx]
            # print(im_label, im_filename, im_data)

            im_label_name = LABELS[im_label]
            im_filename = im_filename.decode()
            im_data = np.array(im_data).reshape((3, 32, 32)).transpose(1, 2, 0)

            # print(im_label_name, im_filename, im_data)
            # # cv2.imshow("im_data", im_data)
            # cv2.imshow("im_data", cv2.resize(im_data, (200, 200)))
            # cv2.waitKey(0)
            if (im_idx + 1) % 1000 == 0:
                print(f"正在写入{file}数据集中，第{im_idx + 1}张图片：{im_filename}")

            sava_img(TEST_DATA_PATH, im_label_name, im_filename, im_data)


if __name__ == '__main__':
    # read_train()
    read_test()
